#!/usr/bin/env python
import rospy
from std_msgs.msg import Bool, String
from geometry_msgs.msg import Pose, PoseStamped
from moveit_commander import MoveGroupCommander
import moveit_commander
import numpy as np
from tf.transformations import quaternion_matrix
import copy
import sys

# 全局变量
left_pose = None  # 接收到的纸巾目标位姿
yogurt_pose = None  # 接收到的酸奶目标位姿
last_valid_pose = None  # 上一次有效的纸巾位姿
last_valid_pose_yogurt = None  # 上一次有效的酸奶位姿
start_move = False  # 是否开始任务
task_state = "stand_by"  # 任务状态

# 连续检测计数器
paper_detection_count = 0  # 纸巾连续检测计数器
yogurt_detection_count = 0  # 酸奶连续检测计数器
DETECTION_THRESHOLD = 5  # 连续检测阈值，例如 5 帧

# MoveIt 组
move_group = None

# 手控制发布器
hand_control_publisher = None

# 纸巾抓取备选位姿
paper_alternative_poses = [
    {
        "orientation": {
            "x": -0.0321,
            "y": 0.1140,
            "z": 0.9236,
            "w": 0.3645
        }
    },
    {
        "orientation": {
            "x": -0.014,
            "y": 0.332,
            "z": 0.912,
            "w": 0.24
        }
    }
]

# 酸奶抓取备选位姿
yogurt_alternative_poses = [
    {
        "orientation": {
            "x": 0.7233,
            "y": -0.2581,
            "z": 0.6007,
            "w": 0.2219
        }
    },
    {
        "orientation": {
            "x": 0.7524,
            "y": -0.1382,
            "z": 0.6390,
            "w": 0.0793
        }
    }
]

# 固定的释放位置
release_pose = PoseStamped()
release_pose.header.frame_id = "base_link"
release_pose.pose.position.x = 0.6  # 固定释放位置的 x 坐标
release_pose.pose.position.y = -0.55  # 固定释放位置的 y 坐标
release_pose.pose.position.z = 0.85  # 固定释放位置的 z 坐标
release_pose.pose.orientation.x = 0.11236370921869202  # 固定释放位置的姿态
release_pose.pose.orientation.y = 0.15079563957693937
release_pose.pose.orientation.z = 0.7905900885524749
release_pose.pose.orientation.w = 0.5827541366801177


def left_pose_callback(msg):
    """
    接收纸巾目标位姿的回调函数。
    """
    global left_pose, last_valid_pose, paper_detection_count  # 将 global 声明放在开头

    # 检查当前 pose 是否有效（避免接收到无效位姿）
    if msg.pose.position.x == -1 or msg.pose.position.y == -1 or msg.pose.position.z == -1:
        paper_detection_count = 0  # 重置计数器
        if last_valid_pose:
            left_pose = copy.deepcopy(last_valid_pose)
        else:
            return
    else:
        if last_valid_pose:
            current_position = np.array([msg.pose.position.x, msg.pose.position.y, msg.pose.position.z])
            last_position = np.array(
                [last_valid_pose.pose.position.x, last_valid_pose.pose.position.y, last_valid_pose.pose.position.z])
            distance = np.linalg.norm(current_position - last_position)

            if distance > 0.10:
                rospy.loginfo(f"Position change exceeds 1cm ({distance:.4f}m), skipping this pose.")
                return
            else:
                left_pose = copy.deepcopy(msg)
                last_valid_pose = copy.deepcopy(msg)
                paper_detection_count += 1  # 增加计数器
        else:
            left_pose = copy.deepcopy(msg)
            last_valid_pose = copy.deepcopy(msg)
            paper_detection_count += 1  # 增加计数器


def yogurt_pose_callback(msg):
    """
    接收酸奶目标位姿的回调函数。
    """
    global yogurt_pose, last_valid_pose_yogurt, yogurt_detection_count  # 将 global 声明放在开头

    # 检查当前 pose 是否有效（避免接收到无效位姿）
    if msg.pose.position.x == -1 or msg.pose.position.y == -1 or msg.pose.position.z == -1:
        yogurt_detection_count = 0  # 重置计数器
        if last_valid_pose_yogurt:
            yogurt_pose = copy.deepcopy(last_valid_pose_yogurt)
        else:
            return
    else:
        if last_valid_pose_yogurt:
            current_position = np.array([msg.pose.position.x, msg.pose.position.y, msg.pose.position.z])
            last_position = np.array(
                [last_valid_pose_yogurt.pose.position.x, last_valid_pose_yogurt.pose.position.y, last_valid_pose_yogurt.pose.position.z])
            distance = np.linalg.norm(current_position - last_position)

            if distance > 0.10:
                rospy.loginfo(f"Position change exceeds 1cm ({distance:.4f}m), skipping this pose.")
                return
            else:
                yogurt_pose = copy.deepcopy(msg)
                last_valid_pose_yogurt = copy.deepcopy(msg)
                yogurt_detection_count += 1  # 增加计数器
        else:
            yogurt_pose = copy.deepcopy(msg)
            last_valid_pose_yogurt = copy.deepcopy(msg)
            yogurt_detection_count += 1  # 增加计数器


def start_move_callback(msg):
    """
    接收开始命令的回调函数。
    """
    global start_move
    if msg.data:
        rospy.loginfo("Received start_move signal, starting task.")
        start_move = True
    else:
        rospy.loginfo("Received stop_move signal, stopping task.")
        start_move = False


def move_to_pose(pose, link_name, is_yogurt=False):
    """
    移动手臂到指定位姿，使用分步规划策略。
    :param pose: Pose 消息，目标位姿
    :param link_name: 要移动的链接名称
    :param is_yogurt: 是否是酸奶任务
    :return: 是否成功
    """
    global move_group
    print("Target pose:", pose)
    print("Current pose:", move_group.get_current_pose().pose)

    current_pose = move_group.get_current_pose().pose

    distance = np.linalg.norm(
        np.array([pose.pose.position.x, pose.pose.position.y, pose.pose.position.z]) -
        np.array([current_pose.position.x, current_pose.position.y, current_pose.position.z])
    )
    print(f"Distance to target pose: {distance}")

    step_size = 0.5
    if distance > step_size:
        num_steps = int(np.ceil(distance / step_size))

        for i in range(num_steps):
            interpolation_factor = (i + 1) / num_steps
            intermediate_pose = Pose()
            intermediate_pose.position.x = current_pose.position.x + (
                    pose.pose.position.x - current_pose.position.x
            ) * interpolation_factor
            intermediate_pose.position.y = current_pose.position.y + (
                    pose.pose.position.y - current_pose.position.y
            ) * interpolation_factor
            intermediate_pose.position.z = current_pose.position.z + (
                    pose.pose.position.z - current_pose.position.z
            ) * interpolation_factor

            intermediate_pose.orientation = pose.pose.orientation

            move_group.set_pose_target(intermediate_pose)
            plan = move_group.plan()

            if plan[0]:
                success = move_group.execute(plan[1], wait=True)
                if not success:
                    rospy.logwarn(f"Failed to move {link_name} to intermediate pose at step {i + 1}.")
                    if try_alternative_orientations(intermediate_pose, link_name, is_yogurt):
                        continue
                    else:
                        return False
            else:
                rospy.logwarn(f"Failed to generate plan for {link_name} at step {i + 1}.")
                if try_alternative_orientations(intermediate_pose, link_name, is_yogurt):
                    continue
                else:
                    return False

        rospy.loginfo(f"Successfully moved {link_name} to the target pose in {num_steps} steps.")
    else:
        move_group.set_pose_target(pose.pose, link_name)
        plan = move_group.plan()

        if plan[0]:
            success = move_group.execute(plan[1], wait=True)
            if success:
                rospy.loginfo(f"Successfully moved {link_name} to the target pose.")
                return True
            else:
                rospy.logwarn(f"Failed to execute plan for {link_name}.")
                return try_alternative_orientations(pose, link_name, is_yogurt)
        else:
            rospy.logwarn(f"Failed to generate plan for {link_name}.")
            return try_alternative_orientations(pose, link_name, is_yogurt)


def try_alternative_orientations(pose, link_name, is_yogurt=False):
    """
    尝试使用备选姿态进行规划。
    :param pose: 目标位姿 (Pose 或 PoseStamped)
    :param link_name: 要移动的链接名称
    :param is_yogurt: 是否是酸奶任务
    :return: 是否成功
    """
    global move_group, paper_alternative_poses, yogurt_alternative_poses

    if isinstance(pose, PoseStamped):
        pose = pose.pose

    # 选择备选姿态列表
    alternative_poses = yogurt_alternative_poses if is_yogurt else paper_alternative_poses

    # 尝试所有备选姿态
    for i, alt_pose in enumerate(alternative_poses):
        alternative_pose = copy.deepcopy(pose)
        alternative_pose.orientation.x = alt_pose["orientation"]["x"]
        alternative_pose.orientation.y = alt_pose["orientation"]["y"]
        alternative_pose.orientation.z = alt_pose["orientation"]["z"]
        alternative_pose.orientation.w = alt_pose["orientation"]["w"]

        move_group.set_pose_target(alternative_pose, link_name)
        plan = move_group.plan()

        if plan[0]:
            success = move_group.execute(plan[1], wait=True)
            if success:
                rospy.loginfo(
                    f"Successfully moved {link_name} to the target pose using alternative orientation {i + 1}.")
                return True
            else:
                rospy.logwarn(f"Failed to execute plan for {link_name} using alternative orientation {i + 1}.")
        else:
            rospy.logwarn(f"Failed to generate plan for {link_name} using alternative orientation {i + 1}.")

    rospy.logwarn("Failed to move to target pose using all alternative orientations.")
    return False


def move_to_stand_by_pose():
    """
    移动手臂到初始位置。
    """
    global move_group
    stand_by_pose = PoseStamped()
    stand_by_pose.header.frame_id = "base_link"
    stand_by_pose.pose.position.x = 0.5
    stand_by_pose.pose.position.y = -0.55
    stand_by_pose.pose.position.z = 0.8
    stand_by_pose.pose.orientation.x = 0.11236370921869202
    stand_by_pose.pose.orientation.y = 0.15079563957693937
    stand_by_pose.pose.orientation.z = 0.7905900885524749
    stand_by_pose.pose.orientation.w = 0.5827541366801177

    return move_to_pose(stand_by_pose, "right_tool_link")


def adjust_pose_for_paper(pose):
    """
    调整目标位姿以适应 paper 任务。
    :param pose: 原始目标位姿 (PoseStamped)
    :return: 调整后的目标位姿 (PoseStamped)
    """
    adjusted_pose = copy.deepcopy(pose)
    adjusted_pose.pose.position.x += 0.06
    adjusted_pose.pose.position.y += -0.10
    adjusted_pose.pose.position.z += 0.04
    adjusted_pose.pose.orientation.x = 0.0927
    adjusted_pose.pose.orientation.y = 0.2553
    adjusted_pose.pose.orientation.z = 0.9106
    adjusted_pose.pose.orientation.w = 0.3113
    return adjusted_pose


def adjust_pose_for_yogurt(pose):
    """
    调整目标位姿以适应 yogurt 任务。
    :param pose: 原始目标位姿 (PoseStamped)
    :return: 调整后的目标位姿 (PoseStamped)
    """
    adjusted_pose = copy.deepcopy(pose)
    adjusted_pose.pose.position.x += 0.06  # 调整位置
    adjusted_pose.pose.position.y += -0.045
    adjusted_pose.pose.position.z += -0.11
    adjusted_pose.pose.orientation.x = 0.62735
    adjusted_pose.pose.orientation.y = -0.2028
    adjusted_pose.pose.orientation.z = 0.7058
    adjusted_pose.pose.orientation.w = 0.2586
    return adjusted_pose


def task_manager():
    """
    任务管理器，根据任务状态执行相应的动作。
    """
    global task_state, start_move, left_pose, yogurt_pose, hand_control_publisher, release_pose
    global paper_detection_count, yogurt_detection_count, DETECTION_THRESHOLD

    rate = rospy.Rate(10)  # 10Hz
    while not rospy.is_shutdown():
        if not start_move:
            rate.sleep()
            continue

        if task_state == "stand_by":
            rospy.loginfo("Task state: stand_by")

            current_pose = move_group.get_current_pose().pose
            stand_by_pose = PoseStamped()
            stand_by_pose.header.frame_id = "base_link"
            stand_by_pose.pose.position.x = 0.5
            stand_by_pose.pose.position.y = -0.5
            stand_by_pose.pose.position.z = 0.8
            stand_by_pose.pose.orientation.x = 0.11236370921869202
            stand_by_pose.pose.orientation.y = 0.15079563957693937
            stand_by_pose.pose.orientation.z = 0.7905900885524749
            stand_by_pose.pose.orientation.w = 0.5827541366801177

            distance = np.linalg.norm(
                np.array([current_pose.position.x, current_pose.position.y, current_pose.position.z]) -
                np.array([stand_by_pose.pose.position.x, stand_by_pose.pose.position.y, stand_by_pose.pose.position.z])
            )

            if distance > 0.1:
                rospy.loginfo("Arm is not in stand_by pose, moving to stand_by pose...")
                if move_to_stand_by_pose():
                    rospy.loginfo("Arm is now in stand_by pose.")
                else:
                    rospy.logwarn("Failed to move to stand_by pose, retrying...")
                    continue

            if left_pose is not None:
                task_state = "move_above_paper"  # 进入移动到纸巾上方的状态
            else:
                print("left_pose", left_pose)
                rospy.loginfo("Waiting for target pose...")

        elif task_state == "move_above_paper":
            rospy.loginfo("Task state: move_above_paper")
            adjusted_pose = adjust_pose_for_paper(left_pose)
            above_pose = copy.deepcopy(adjusted_pose)
            above_pose.pose.position.z += 0.10  # 移动到纸巾上方 10 厘米

            if move_to_pose(above_pose, "right_tool_link"):
                task_state = "move_to_paper"  # 移动到纸巾上方后进入抓取状态
                rospy.sleep(1)  # 暂停一秒钟
            else:
                rospy.logwarn("Failed to move above paper, retrying...")

        elif task_state == "move_to_paper":
            rospy.loginfo("Task state: move_to_paper")
            adjusted_pose = adjust_pose_for_paper(left_pose)
            if move_to_pose(adjusted_pose, "right_tool_link"):
                task_state = "grasp"
                rospy.sleep(1)  # 暂停一秒钟
            else:
                rospy.logwarn("Failed to move to paper pose, retrying...")

        elif task_state == "grasp":
            rospy.loginfo("Task state: grasp")
            msg_hand_control = String()
            msg_hand_control.data = "paper_grasp"  # 发布抓取纸巾的信息
            hand_control_publisher.publish(msg_hand_control)
            rospy.loginfo("Sent hand close command for paper.")
            rospy.sleep(2)
            task_state = "move_up_after_grasp"  # 抓取完成后进入向上移动状态
            rospy.sleep(3)  # 暂停一秒钟

        elif task_state == "move_up_after_grasp":
            rospy.loginfo("Task state: move_up_after_grasp")
            current_pose = move_group.get_current_pose().pose
            target_pose = copy.deepcopy(current_pose)
            target_pose.position.z += 0.10  # 向上移动 10 厘米

            target_pose_stamped = PoseStamped()
            target_pose_stamped.header.frame_id = "base_link"
            target_pose_stamped.pose = target_pose

            if move_to_pose(target_pose_stamped, "right_tool_link"):
                task_state = "move_to_release"  # 向上移动完成后进入释放状态
                rospy.sleep(1)  # 暂停一秒钟
            else:
                rospy.logwarn("Failed to move up after grasp, retrying...")

        elif task_state == "move_to_release":
            rospy.loginfo("Task state: move_to_release")
            if move_to_pose(release_pose, "right_tool_link"):
                task_state = "release"
                rospy.sleep(1)  # 暂停一秒钟
            else:
                rospy.logwarn("Failed to move to release pose, retrying...")

        elif task_state == "release":
            rospy.loginfo("Task state: release")
            msg_hand_control = String()
            msg_hand_control.data = "release"  # 发布释放纸巾的信息
            hand_control_publisher.publish(msg_hand_control)
            rospy.loginfo("Sent hand open command for paper.")
            rospy.sleep(2)
            task_state = "return_home_after_paper"  # 扔完纸巾后回到初始位置
            rospy.sleep(1)  # 暂停一秒钟

        elif task_state == "return_home_after_paper":
            rospy.loginfo("Task state: return_home_after_paper")
            if move_to_stand_by_pose():
                rospy.loginfo("Returned to stand_by pose after paper task.")

                # 重置 last_valid_pose 和计数器
                global last_valid_pose, paper_detection_count
                last_valid_pose = None
                paper_detection_count = 0  # 重置计数器

                # 检查视野中是否还有酸奶
                if yogurt_detection_count >= DETECTION_THRESHOLD:
                    rospy.loginfo("Detected remaining yogurt, executing yogurt task...")
                    task_state = "move_above_yogurt"  # 进入酸奶任务
                    start_move = True  # 重新开始任务
                else:
                    rospy.loginfo("No remaining yogurt detected, returning to stand_by state.")
                    task_state = "stand_by"  # 返回初始状态
                    start_move = False  # 任务完成，等待下一次开始命令

                rospy.sleep(1)  # 暂停一秒钟
            else:
                rospy.logwarn("Failed to return to stand_by pose after paper task, retrying...")

        elif task_state == "move_above_yogurt":
            rospy.loginfo("Task state: move_above_yogurt")
            if yogurt_pose is None:
                rospy.loginfo("Waiting for yogurt pose...")
                rate.sleep()
                continue

            adjusted_yogurt_pose = adjust_pose_for_yogurt(yogurt_pose)
            above_pose = copy.deepcopy(adjusted_yogurt_pose)
            above_pose.pose.position.x += -0.01
            above_pose.pose.position.y += -0.10
            above_pose.pose.position.z += 0.10  # 移动到酸奶上方 10 厘米

            if move_to_pose(above_pose, "right_tool_link", is_yogurt=True):
                task_state = "move_to_yogurt"  # 移动到酸奶上方后进入抓取状态
                rospy.sleep(1)  # 暂停一秒钟
            else:
                rospy.logwarn("Failed to move above yogurt, retrying...")

        elif task_state == "move_to_yogurt":
            rospy.loginfo("Task state: move_to_yogurt")
            adjusted_yogurt_pose = adjust_pose_for_yogurt(yogurt_pose)
            if move_to_pose(adjusted_yogurt_pose, "right_tool_link", is_yogurt=True):
                task_state = "grasp_yogurt"
                rospy.sleep(1)  # 暂停一秒钟
            else:
                rospy.logwarn("Failed to move to yogurt pose, retrying...")

        elif task_state == "grasp_yogurt":
            rospy.loginfo("Task state: grasp_yogurt")
            msg_hand_control = String()
            msg_hand_control.data = "yogurt_grasp"  # 发布抓取酸奶的信息
            hand_control_publisher.publish(msg_hand_control)
            rospy.loginfo("Sent hand close command for yogurt.")
            rospy.sleep(2)
            task_state = "move_up_after_grasp_yogurt"  # 抓取酸奶后进入向上移动状态
            rospy.sleep(1)  # 暂停一秒钟

        elif task_state == "move_up_after_grasp_yogurt":
            rospy.loginfo("Task state: move_up_after_grasp_yogurt")
            current_pose = move_group.get_current_pose().pose
            target_pose = copy.deepcopy(current_pose)
            target_pose.position.z += 0.10  # 向上移动 10 厘米

            target_pose_stamped = PoseStamped()
            target_pose_stamped.header.frame_id = "base_link"
            target_pose_stamped.pose = target_pose

            if move_to_pose(target_pose_stamped, "right_tool_link", is_yogurt=True):
                task_state = "move_to_yogurt_release"  # 向上移动完成后进入释放状态
                rospy.sleep(1)  # 暂停一秒钟
            else:
                rospy.logwarn("Failed to move up after yogurt grasp, retrying...")

        elif task_state == "move_to_yogurt_release":
            rospy.loginfo("Task state: move_to_yogurt_release")
            if move_to_pose(release_pose, "right_tool_link", is_yogurt=True):
                task_state = "release_yogurt"
                rospy.sleep(1)  # 暂停一秒钟
            else:
                rospy.logwarn("Failed to move to yogurt release pose, retrying...")

        elif task_state == "release_yogurt":
            rospy.loginfo("Task state: release_yogurt")
            msg_hand_control = String()
            msg_hand_control.data = "release"  # 发布释放酸奶的信息
            hand_control_publisher.publish(msg_hand_control)
            rospy.loginfo("Sent hand open command for yogurt.")
            rospy.sleep(2)
            task_state = "return_home_after_yogurt"  # 释放酸奶后回到初始位置
            rospy.sleep(1)  # 暂停一秒钟

        elif task_state == "return_home_after_yogurt":
            rospy.loginfo("Task state: return_home_after_yogurt")
            if move_to_stand_by_pose():
                rospy.loginfo("Task completed, returning to stand_by state.")

                # 重置 last_valid_pose_yogurt 和计数器
                global last_valid_pose_yogurt, yogurt_detection_count
                last_valid_pose_yogurt = None
                yogurt_detection_count = 0  # 重置计数器

                # 检查视野中是否还有纸巾
                if paper_detection_count >= DETECTION_THRESHOLD:
                    rospy.loginfo("Detected remaining paper, executing paper task...")
                    task_state = "move_above_paper"  # 进入纸巾任务
                    start_move = True  # 重新开始任务
                else:
                    rospy.loginfo("No remaining paper detected, task fully completed.")
                    task_state = "stand_by"  # 返回初始状态
                    start_move = False  # 任务完成，等待下一次开始命令

                rospy.sleep(1)  # 暂停一秒钟
            else:
                rospy.logwarn("Failed to return to stand_by pose after yogurt task, retrying...")




if __name__ == "__main__":
    rospy.init_node('move_arm_task_manager', anonymous=True)

    moveit_commander.roscpp_initialize(sys.argv)
    move_group = MoveGroupCommander("right_arm")
    move_group.set_end_effector_link("right_tool_link")
    move_group.set_planning_time(0.2)
    move_group.set_num_planning_attempts(10)
    move_group.set_max_velocity_scaling_factor(0.3)
    move_group.set_max_acceleration_scaling_factor(0.3)
    move_group.set_goal_position_tolerance(0.01)
    move_group.set_goal_orientation_tolerance(0.3)
    move_group.set_planner_id("PRM")

    hand_control_publisher = rospy.Publisher("/hand_control_topic", String, queue_size=1)

    # 订阅纸巾和酸奶的位姿
    rospy.Subscriber('/paper_pose', PoseStamped, left_pose_callback, queue_size=1)
    rospy.Subscriber('/yoger_pose', PoseStamped, yogurt_pose_callback, queue_size=1)
    rospy.Subscriber('/start_move', Bool, start_move_callback, queue_size=1)

    task_manager()

    rospy.spin()